#!/usr/bin/env python3
import os
import sys
import json
import fcntl
from pathlib import Path
from dotenv import load_dotenv
import requests
from colorama import Fore, Style
import asyncio
import aiohttp
from tqdm.asyncio import tqdm_asyncio
import mimetypes

def get_script_dir() -> Path:
    """始终返回脚本所在目录的绝对路径"""
    script_path = Path(__file__).resolve()
    return script_path.parent

SCRIPT_DIR = get_script_dir()

HISTORY_FILE = SCRIPT_DIR / ".deepseek_token_history"
LOCK_FILE = SCRIPT_DIR / ".deepseek_token.lock"
DEEPSEEK_CLI_VERSION = "0.0.0.202504201829"

# 新增异步进度条装饰器
def async_progress(func):
    async def wrapper(*args, **kwargs):
        with tqdm_asyncio(total=4, desc=f"{Fore.RED}DeepSeek处理进度{Style.RESET_ALL}",
                        bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}") as pbar:
            async def update_progress():
                stages = ["初始化请求", "建立连接", "传输数据", "解析响应"]
                for i in range(4):
                    await asyncio.sleep(0.3)
                    pbar.set_postfix_str(f"{Fore.YELLOW}{stages[i]}{Style.RESET_ALL}")
                    pbar.update(1)
            task = asyncio.create_task(update_progress())
            result = await func(*args, **kwargs)
            task.cancel()
            return result
    return wrapper

def acquire_lock():
    """文件锁防止并发写入"""
    lock = open(LOCK_FILE, 'w')
    fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
    return lock

def release_lock(lock):
    fcntl.flock(lock.fileno(), fcntl.LOCK_UN)
    lock.close()

def read_history():
    """读取历史记录，包含容错机制"""
    try:
        with open(HISTORY_FILE, 'r') as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {"total_tokens": 0, "history": []}

def write_history(data):
    """原子化写入历史记录"""
    lock = acquire_lock()
    try:
        with open(HISTORY_FILE, 'w') as f:
            json.dump(data, f, indent=2)
    finally:
        release_lock(lock)

def update_token_history(usage_data):
    """更新token统计"""
    history = read_history()

    # 本次消耗
    current_usage = {
        "prompt_tokens": usage_data.get("prompt_tokens", 0),
        "completion_tokens": usage_data.get("completion_tokens", 0),
        "total_tokens": usage_data.get("total_tokens", 0)
    }

    # 更新历史记录
    history["history"].append(current_usage)
    history["total_tokens"] += current_usage["total_tokens"]

    write_history(history)
    return current_usage, history["total_tokens"]

def format_token_output(current, total):
    return (
        f"{Fore.YELLOW}Token统计{Style.RESET_ALL}\n"
        f"├─ 本次Prompt: {Fore.CYAN}{current['prompt_tokens']}{Style.RESET_ALL}\n"
        f"├─ 本次Completion: {Fore.CYAN}{current['completion_tokens']}{Style.RESET_ALL}\n"
        f"├─ 本次Total: {Fore.GREEN}{current['total_tokens']}{Style.RESET_ALL}\n"
        f"└─ 历史累计: {Fore.MAGENTA}{total}{Style.RESET_ALL}"
    )

# 异步化核心请求函数
async def async_chat_with_deepseek(prompt):
    headers = {
        "Authorization": f"Bearer {os.getenv('DEEPSEEK_API_KEY')}",
        "Content-Type": "application/json"
    }
    payload = {
        "messages": [{"role": "user", "content": prompt}],
        "model": "deepseek-chat",
        "temperature": 0.7,
        "stream": False
    }

    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30)) as session:
            async with session.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers=headers,
                json=payload
            ) as response:
                data = await response.json()
                usage = data.get("usage", {})
                # 更新token统计
                current_usage, total = update_token_history(usage)
                return (
                    f"\n{Fore.YELLOW}DeepSeek:{Style.RESET_ALL}\n"
                    f"{data['choices'][0]['message']['content']}\n\n"
                    f"{format_token_output(current_usage, total)}"
                )
    except Exception as e:
        return f"{Fore.RED}Error: {str(e)}{Style.RESET_ALL}"

# 组合式进度条
class DualProgress:
    def __init__(self):
        self.spinner = None
        self.bar = None
    async def start(self):
        self.spinner = asyncio.create_task(self.spin())
        self.bar = tqdm_asyncio(total=100, desc=f"{Fore.CYAN}传输进度{Style.RESET_ALL}")
    async def spin(self):
        symbols = '🕐🕑🕒🕓🕔🕕🕖🕗🕘🕙🕚🕛'
        while True:
            for s in symbols:
                print(f"\r{s} {Fore.MAGENTA}等待API响应...{Style.RESET_ALL}", end='', flush=True)
                await asyncio.sleep(0.3)
    async def update(self, value):
        if self.bar.n + value < 100:
            self.bar.update(value)
    async def stop(self):
        self.spinner.cancel()
        self.bar.close()
        print('\r' + ' '*50 + '\r', end='')

# 集成进度条的主函数
@async_progress
async def enhanced_chat(prompt):
    progress = DualProgress()
    await progress.start()
    try:
        # 模拟分阶段进度
        await asyncio.gather(
            progress.update(30),
            asyncio.sleep(0.5),
            progress.update(30),
            asyncio.sleep(0.5),
            progress.update(40)
        )
        return await async_chat_with_deepseek(prompt)
    finally:
        await progress.stop()

def analyze_code_file(file_path: Path) -> str:
    """智能识别代码文件类型并生成分析提示"""
    try:
        content = file_path.read_text(encoding='utf-8')
    except UnicodeDecodeError:
        return f"{Fore.RED}错误：文件编码不支持（请使用UTF-8编码）{Style.RESET_ALL}"

    # 根据文件后缀生成类型描述
    file_type_map = {
        '.c': 'C语言程序',
        '.py': 'Python脚本',
        '.java': 'Java程序',
        '.js': 'JavaScript代码',
        '.html': 'HTML文档',
        '.md': 'Markdown文档'
    }
    file_type = file_type_map.get(file_path.suffix.lower(), "代码文件")

    return (
        f"请分析以下{file_type}的功能和实现逻辑：\n"
        f"文件路径：{file_path}\n"
        f"文件内容：\n```\n{content}\n```\n"
        "请按照以下结构进行专业分析：\n"
        "1. 核心功能描述\n"
        "2. 关键算法/实现解析\n"
        "3. 潜在改进建议\n"
        "4. 安全风险提示"
    )

def generate_prompt(args: list) -> str:
    """生成智能分析提示"""
    # 分离文件和普通参数
    files = []
    other_args = []

    for arg in args:
        file_path = Path(arg)
        if file_path.exists() and file_path.is_file():
            files.append(file_path)
        else:
            other_args.append(arg)

    # 优先处理第一个代码文件
    if files:
        main_file = files[0]
        prompt = analyze_code_file(main_file)

        # 添加附加问题
        if other_args:
            prompt += f"\n\n用户附加问题：{' '.join(other_args)}"

        # 添加多文件标记
        if len(files) > 1:
            prompt += f"\n\n注：检测到共{len(files)}个相关文件，可指定文件进行深入分析"

        return prompt

    return ' '.join(args)

def show_help():
    print(f"{Fore.RED}使用方式：")
    print(f"{Fore.RED}Usage: deepseek <query>{Style.RESET_ALL}")
    print(f"  代码分析：deepseek 文件路径 [附加问题]{Style.RESET_ALL}")
    print(f"  普通问答模式：deepseek \"你的问题\"{Style.RESET_ALL}")
    print(f"  查看帮助：deepseek -h")
    print(f"  查看版本：deepseek -v")

def show_version():
    print(DEEPSEEK_CLI_VERSION)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        show_help()
        sys.exit(1)
    query = " ".join(sys.argv[1:])
    if query == "-h":
        show_help()
        sys.exit(1)
    if query == "-v":
        show_version()
        sys.exit(1)
    load_dotenv()
    prompt = generate_prompt(query)
    print(f"\n{Fore.RED}请稍候，数据请求中......{Style.RESET_ALL}")
    # 获取事件循环
    if sys.platform == 'win32':
        loop = asyncio.ProactorEventLoop()
        asyncio.set_event_loop(loop)
    else:
        loop = asyncio.get_event_loop()
    response = loop.run_until_complete(enhanced_chat(prompt))
    print(response)
